import re

import requests
from fastapi import APIRouter, UploadFile, File, HTTPException, Depends, Query
from typing import Optional, Dict, Any, List
from pydantic import BaseModel
import os
import tempfile
from pathlib import Path

# 导入下方工具类
from tools.rag_file_processor import RAGFileProcessor, create_rag_processor
from tools.small_chat_tools import Knowledge

# 初始化路由实例
router = APIRouter(
    prefix="/knowledge",
    tags=["知识库操作"]
)



# 核心依赖与工具类
def get_knowledge_client() -> Knowledge:
    """依赖注入：获取知识库客户端"""
    try:
        knowledge = Knowledge()
        knowledge.get_document_count()
        return knowledge
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"知识库客户端初始化失败: {str(e)}"
        )

class HtmlUrlRequest(BaseModel):
    """HTML网页检索请求模型"""
    url: str
    metadata: Optional[Dict[str, Any]] = None



# 请求模型（目录处理）
class DirectoryProcessRequest(BaseModel):
    directory_path: str
    recursive: bool = False
    metadata: Optional[Dict[str, Any]] = None



# 多格式文件上传（支持PDF/DOCX/MD/TXT）
@router.post("/upload-rag-file", summary="上传多格式文件至RAG知识库")
def upload_rag_file(
        file: UploadFile = File(..., description="支持的文件格式: PDF, DOCX, MD, TXT"),
        metadata: Optional[Dict[str, Any]] = None,
        chunk_size: Optional[int] = Query(None, description="文本块大小"),
        chunk_overlap: Optional[int] = Query(None, description="文本块重叠大小"),
        knowledge: Knowledge = Depends(get_knowledge_client)
):
    """上传并处理多格式文件，存储为RAG文本块"""
    processor = RAGFileProcessor(knowledge=knowledge)
    try:
        success, result = processor.process_uploaded_file(
            file=file,
            metadata=metadata
        )
        if success:
            return {
                "code": 200,
                "message": "RAG文件处理成功",
                "data": result
            }
        raise HTTPException(
            status_code=400,
            detail=result.get("error", "文件处理失败")
        )
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"处理RAG文件时发生错误: {str(e)}"
        )


# 文本文件上传（TXT专用，大粒度分割）
@router.post("/upload-file", summary="上传文本文件至知识库（TXT专用）")
def upload_file_to_knowledge(
        file: UploadFile = File(...),
        metadata: Optional[Dict[str, Any]] = None,
        encoding: str = "utf-8",
        knowledge: Knowledge = Depends(get_knowledge_client)
):
    """上传TXT文件，按段落分割后存储"""
    if not file.filename.endswith(".txt"):
        raise HTTPException(
            status_code=400,
            detail="仅支持.txt格式"
        )

    processor = RAGFileProcessor(knowledge=knowledge)
    try:
        # 临时存储TXT文件，再用工具类处理
        with tempfile.NamedTemporaryFile(delete=False, suffix=".txt") as temp_file:
            content = file.file.read().decode(encoding)
            temp_file.write(content.encode(encoding))
            temp_file_path = temp_file.name

        # 按段落分割（工具类已有process_and_store_document逻辑，这里复用）
        success, result = processor.process_and_store_document(
            file_path=temp_file_path,
            metadata=metadata
        )

        if success:
            return {
                "code": 200,
                "message": "TXT文件录入成功",
                "data": {
                    "filename": file.filename,
                    "total_paragraphs": result.get("raw_pages", 0),
                    "valid_count": result.get("inserted_count", 0),
                    "current_total": knowledge.get_document_count(),
                    "metadata": metadata
                }
            }
        raise HTTPException(
            status_code=400,
            detail=result.get("error", "TXT文件处理失败")
        )
    except UnicodeDecodeError:
        raise HTTPException(
            status_code=400,
            detail=f"编码错误，当前编码{encoding}"
        )
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"处理TXT文件时发生错误: {str(e)}"
        )
    finally:
        if 'temp_file_path' in locals() and os.path.exists(temp_file_path):
            os.unlink(temp_file_path)


# 处理服务器本地目录中的文件
@router.post("/process-directory", summary="批量处理服务器目录中的文件")
def process_directory(
        request: DirectoryProcessRequest,
        knowledge: Knowledge = Depends(get_knowledge_client)
):
    """批量处理本地目录中所有支持格式的文件"""
    if not os.path.isdir(request.directory_path):
        raise HTTPException(
            status_code=404,
            detail=f"目录不存在: {request.directory_path}"
        )

    processor = RAGFileProcessor(knowledge=knowledge)
    try:
        result = processor.process_and_store_directory(
            dir_path=request.directory_path,
            recursive=request.recursive,
            metadata=request.metadata
        )
        return {
            "code": 200,
            "message": "目录批量处理完成",
            "data": result
        }
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"批量处理目录失败: {str(e)}"
        )


# 根据ID删除知识库文档
@router.delete("/delete-by-ids", summary="根据ID删除知识库文档")
def delete_documents_by_ids(
        ids: List[str],
        knowledge: Knowledge = Depends(get_knowledge_client)
):
    """根据ID列表删除文档"""
    if not ids:
        raise HTTPException(
            status_code=400,
            detail="未提供文档ID"
        )

    try:
        before_count = knowledge.get_document_count()
        success_count = 0
        failed_ids = []

        for doc_id in ids:
            if knowledge.delete_document(doc_id):
                success_count += 1
            else:
                failed_ids.append(doc_id)

        after_count = knowledge.get_document_count()

        return {
            "code": 200,
            "message": f"成功删除 {success_count} 个文档，失败 {len(failed_ids)} 个文档",
            "data": {
                "total_ids": len(ids),
                "success_count": success_count,
                "failed_ids": failed_ids,
                "before_count": before_count,
                "after_count": after_count,
                "deleted_count": before_count - after_count
            }
        }
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"删除文档失败: {str(e)}"
        )


# 查看知识库状态
@router.get("/status", summary="查看知识库当前状态")
def get_knowledge_status(
        knowledge: Knowledge = Depends(get_knowledge_client)
):
    """获取知识库基本信息及文档总数"""
    try:
        return {
            "code": 200,
            "message": "获取知识库状态成功",
            "data": {
                "database": knowledge.milvus.FIXED_DB_NAME,
                "collection": knowledge.milvus.FIXED_COLLECTION_NAME,
                "document_count": knowledge.get_document_count()
            }
        }
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"获取知识库状态失败: {str(e)}"
        )


# HTML网页检索接口
@router.post("/fetch-html-url", summary="从URL获取HTML网页并存储到知识库")
def fetch_and_store_html(
        request: HtmlUrlRequest,
        knowledge: Knowledge = Depends(get_knowledge_client)
):
    """
    从指定URL获取HTML网页内容，清洗后分割存储到知识库
    - 自动过滤script、style等无效标签
    - 支持自定义元数据（如分类、来源等）
    """
    # 基础URL校验
    if not re.match(r'^https?://', request.url):
        raise HTTPException(
            status_code=400,
            detail="URL格式无效，请使用http://或https://开头的URL"
        )

    processor = RAGFileProcessor(knowledge=knowledge)
    try:
        success, result = processor.fetch_and_process_html(
            url=request.url,
            metadata=request.metadata
        )
        if success:
            return {
                "code": 200,
                "message": "HTML网页处理成功",
                "data": result
            }
        raise HTTPException(
            status_code=400,
            detail=result.get("error", "HTML网页处理失败")
        )
    except requests.exceptions.Timeout:
        raise HTTPException(
            status_code=408,
            detail="请求URL超时，请检查网络或URL有效性"
        )
    except requests.exceptions.RequestException as e:
        raise HTTPException(
            status_code=400,
            detail=f"网络请求失败: {str(e)}"
        )
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"处理HTML网页时发生错误: {str(e)}"
        )


# HTML文件上传
@router.post("/upload-rag-file", summary="上传多格式文件至RAG知识库（支持HTML）")
def upload_rag_file(
        file: UploadFile = File(..., description="支持的文件格式: PDF, DOCX, MD, TXT, HTML"),
        metadata: Optional[Dict[str, Any]] = None,
        chunk_size: Optional[int] = Query(None, description="文本块大小"),
        chunk_overlap: Optional[int] = Query(None, description="文本块重叠大小"),
        knowledge: Knowledge = Depends(get_knowledge_client)
):
    """上传并处理多格式文件（新增HTML支持），存储为RAG文本块"""
    # 检查文件格式是否支持（明确提示HTML支持）
    supported_exts = RAGFileProcessor().supported_extensions
    file_ext = Path(file.filename).suffix.lower()
    if file_ext not in supported_exts:
        raise HTTPException(
            status_code=400,
            detail=f"不支持的文件格式: {file_ext}，支持格式: {list(supported_exts.keys())}"
        )

    processor = RAGFileProcessor(knowledge=knowledge)
    try:
        success, result = processor.process_uploaded_file(
            file=file,
            metadata=metadata
        )
        if success:
            return {
                "code": 200,
                "message": "文件处理成功",
                "data": result
            }
        raise HTTPException(
            status_code=400,
            detail=result.get("error", "文件处理失败")
        )
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"处理文件时发生错误: {str(e)}"
        )


# 更新支持的文件格式说明接口
@router.get("/supported-formats", summary="获取支持的文件格式")
def get_supported_formats(knowledge: Knowledge = Depends(get_knowledge_client)):
    """返回支持的文件格式及说明（新增HTML说明）"""
    processor = RAGFileProcessor(knowledge=knowledge)
    return {
        "code": 200,
        "message": "获取支持的文件格式成功",
        "data": {
            "supported_formats": processor.supported_extensions,
            "special_notes": {
                "txt": "推荐使用 /upload-file 接口",
                "html": "支持本地HTML文件上传（/upload-rag-file）和在线URL检索（/fetch-html-url）"
            }
        }
    }


# 获取所有文档ID
@router.get("/all-ids", summary="获取知识库中所有文档ID")
def get_all_document_ids(
        limit: int = 1000,
        knowledge: Knowledge = Depends(get_knowledge_client)
):
    """获取所有文档ID列表，辅助删除操作"""
    try:
        from pymilvus import Collection

        if not knowledge.milvus.client.has_collection(knowledge.fixed_collection):
            knowledge.milvus.create_collection(
                collection_name=knowledge.fixed_collection,
                dimension=knowledge.milvus.FIXED_VECTOR_DIM
            )

        collection = Collection(knowledge.fixed_collection)
        collection.load()

        result = collection.query(
            expr="id >= 0",
            output_fields=["id"],
            limit=limit
        )

        all_ids = [str(item["id"]) for item in result]

        return {
            "code": 200,
            "message": f"成功获取 {len(all_ids)} 个文档ID",
            "data": {
                "ids": all_ids,
                "total_count": len(all_ids),
                "limit": limit
            }
        }
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"获取文档ID列表失败: {str(e)}"
        )


# 获取支持的文件格式列表
@router.get("/supported-formats", summary="获取支持的文件格式")
def get_supported_formats(knowledge: Knowledge = Depends(get_knowledge_client)):
    """返回支持的文件格式及说明"""
    processor = RAGFileProcessor(knowledge=knowledge)
    return {
        "code": 200,
        "message": "获取支持的文件格式成功",
        "data": {
            "supported_formats": processor.supported_extensions,
            "txt_special_note": "TXT格式推荐使用 /upload-file 接口，其他格式使用 /upload-rag-file 接口"
        }
    }